// websocket 基础版
package serve

import (
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

func ServeWebSocket(c *gin.Context) {
	var upgrader = websocket.Upgrader{
		HandshakeTimeout: time.Second * 10,
		ReadBufferSize:   1024,
		WriteBufferSize:  1024,
		CheckOrigin: func(r *http.Request) bool {
			// 在这里进行权限验证，返回true表示验证通过，允许连接
			// 可以根据需要实现自己的验证逻辑
			log.Println("升级协议", r.Header["User-Agent"])
			return true
		},
	}

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	userId := c.Query("uid")
	log.Println("用户id:", userId)

	// 将连接加入连接池
	addSession(userId, conn)

	// 监听客户端发送的消息
	for {
		// 读取客户端发送的消息
		messageType, message, err := conn.ReadMessage()
		if err != nil {
			log.Println(err)
			break
		}

		fmt.Printf("获取客户端发送的消息: %s, messageType: %d\n", string(message), messageType)

		// 发送消息到客户端
		switch messageType {
		case websocket.TextMessage:
			fmt.Printf("处理文本消息, %s\n", string(message))
			// 处理心跳请求
			if string(message) == "ping" {
				message = []byte("pong")
			}
			conn.WriteMessage(websocket.TextMessage, message)
		case websocket.BinaryMessage:
			fmt.Println("处理二进制消息")
		case websocket.CloseMessage:
			fmt.Println("关闭websocket连接")
			return
		case websocket.PingMessage:
			fmt.Println("处理ping消息")
			conn.WriteMessage(websocket.PongMessage, []byte("ping"))
		case websocket.PongMessage:
			fmt.Println("处理pong消息")
			conn.WriteMessage(websocket.PongMessage, []byte("pong"))
		default:
			fmt.Printf("未知消息类型: %d\n", messageType)
			return
		}
	}

	// 连接断开时，从连接池中移除
	removeSession(userId)
}

var connects = make(map[string]*websocket.Conn)
var lock sync.RWMutex

func addSession(id string, conn *websocket.Conn) {
	lock.Lock()
	connects[id] = conn
	defer lock.Unlock()
}

func getSession(id string) *websocket.Conn {
	lock.RLock()
	defer lock.RUnlock()
	return connects[id]
}

func removeSession(id string) {
	lock.Lock()
	delete(connects, id)
	defer lock.Unlock()
}
